import torch
import torchvision
from torch import nn
from torch.utils.tensorboard import SummaryWriter


class MyNN(torch.nn.Module):
    def __init__(self):
        super(MyNN, self).__init__()
        # self.conv1 = nn.Conv2d(3,32,5,padding=2)
        # self.maxpool1 = nn.MaxPool2d(2)
        # self.conv2 = nn.Conv2d(32,32,5,padding=2)
        # self.maxpool2 = nn.MaxPool2d(2)
        # self.conv3 = nn.Conv2d(32,64,5,padding=2)
        # self.maxpool3 = nn.MaxPool2d(2)
        # self.flatten = nn.Flatten()
        # self.linear1 = nn.Linear(1024,64)
        # self.linear2 = nn.Linear(64,10)

        self.model1 = nn.Sequential(nn.Conv2d(3,32,5,padding=2),
                                    nn.MaxPool2d(2),
                                    nn.Conv2d(32,32,5,padding=2),
                                    nn.MaxPool2d(2),
                                    nn.Conv2d(32,64,5,padding=2),
                                    nn.MaxPool2d(2),
                                    nn.Flatten(),
                                    nn.Linear(1024,64),
                                    nn.Linear(64,10))

    def forward(self,input):
        # input = self.conv1(input)
        # input = self.maxpool1(input)
        # input = self.conv2(input)
        # input = self.maxpool2(input)
        # input = self.conv3(input)
        # input = self.maxpool3(input)
        # input = self.flatten(input)
        # input = self.linear1(input)
        # output = self.linear2(input)

        output = self.model1(input)

        return output

my_nn = MyNN()
print(my_nn)
input = torch.zeros((64,3,32,32))
output = my_nn(input)
print(output.shape)

writer = SummaryWriter("logs_seq")
writer.add_graph(my_nn,input)
writer.close()